//========================================================================================
// SuffixTree.hh
// ---------------------------------------------------------------------------------------
// $Id: SuffixTree.hh,v 1.2 2000/03/22 14:16:22 kdurbin Exp kdurbin $
// ---------------------------------------------------------------------------------------
// $Log: SuffixTree.hh,v $
// Revision 1.2  2000/03/22 14:16:22  kdurbin
// Maintenence check-in: ensure latest .cc has correct .hh
//
// Revision 1.1  2000/03/08 00:52:29  kdurbin
// Initial revision
//
//========================================================================================
//.......1.........2.........3.........4.........5.........6.........7.........8.........9
//

#ifndef _SUFFIX_TREE
#define _SUFFIX_TREE

#include "CBT.hh" 

const int CHARACTER_SET_SIZE = 6;


//-------------------------------------------------------------------
// * SuffixTreeSuffix
//------------------------------------------------------------------
// When a new tree is added to the table, we step through all the currently defined suffixes 
// from the active point to the end point.  This structure defines a Suffix by its final 
// character. In the canonical representation, we define that last character by starting at 
// a node in the tree, and following a string of characters, represented by first_char_index 
// and last_char_index.  The two indices point into the input string.  Note that if a suffix
// ends at a node, there are no additional characters needed to characterize its last character 
// position. When this is the case, we say the node is Explicit, and set first_char_index > 
// last_char_index to flag that.
//
class SuffixTreeSuffix {
public :
	
  long mOriginNode;

  // This is the index of the character leading out of mOriginNode that contains 
  // this suffix.  In effect, it specifies the edge that contains the suffix. 
  long mFirstCharIdx;
  long mLastCharIdx;
  
  void Print(){
	 cout<<"Node="<<mOriginNode;
	 cout<<" 1stChar="<<mFirstCharIdx<<" LastChar="<<mLastCharIdx<<endl;
  };

  SuffixTreeSuffix( long node, long start, long stop )
	 : mOriginNode( node ),
		mFirstCharIdx( start ),
		mLastCharIdx( stop ){};
  
  inline bool Explicit(){ return(mFirstCharIdx > mLastCharIdx); }
  inline bool Implicit(){ return(mLastCharIdx >= mFirstCharIdx);}
};


//-------------------------------------------------------------------
// * SuffixTreeEdge
//------------------------------------------------------------------
// The suffix tree is made up of edges connecting nodes. Each edge represents a string 
// of characters starting at first_char_index and ending at last_char_index. Edges can be 
// inserted and removed from a hash table, based on the Hash() function defined here.  The 
// hash table indicates an unused slot by setting the mFirstCharIdx to -1.
//
class SuffixTreeEdge {
public :  
	long mFirstCharIdx;
	long mLastCharIdx;
	long mEndNode;
	//long mStartNode;  // Not actually used in algorithm. 
	
	SuffixTreeEdge(){Invalidate();};

  void Print(){
	 cout<<"FirstCharIdx="<<mFirstCharIdx;
	 cout<<" LastCharIdx="<<mLastCharIdx;
	 cout<<" EndNode="<<mEndNode<<endl;
  };
              
 	void Invalidate(){mFirstCharIdx = -1;mLastCharIdx = -1;mEndNode=-1;};
};

//-------------------------------------------------------------------
// * SuffixTreeNode
//------------------------------------------------------------------
// The only information contained in a node is the suffix link. Each suffix in the tree 
// that ends at a particular node can find the next smaller suffix by following the 
// suffix_node link to a new node.  Nodes are stored in a simple array.
//
class SuffixTreeNode {
public :

  SuffixTreeNode(){ mSuffixNode = -1;
  			 mLeafIdx = -1;};

  void Initialize(){
	 for(int i = 0;i<CHARACTER_SET_SIZE;i++){
		mEdges[i].Invalidate();
	 }
  };
		
  void Print(){
	 cout<<"SuffixNode="<<mSuffixNode<<" LeafIdx="<<mLeafIdx<<endl;
	 for(int i = 0;i<CHARACTER_SET_SIZE;i++){
		mEdges[i].Print();
	 }
  }


  SuffixTreeEdge mEdges[CHARACTER_SET_SIZE];
  long mSuffixNode;    
  long mLeafIdx;

};


//-------------------------------------------------------------------
// * STreeSearchStack
//------------------------------------------------------------------
// Very simple stack for the DFS search.  Since we know the worst case bound on the DFS, 
// space can be pre-allocated.  The stack can be quite deep, so simply doing recursive
// function calls and using the system stack can result in stack overflows.   This stack
// is also faster.  
//
class STreeSearchStack{
public:

  long *mpStackData;
  bool *mpVisited;
  int mStackPtr;
  int mSize;
  int mMaxVisited;

  STreeSearchStack(int Size){
	 mpStackData = new long[Size];
	 mpVisited = new bool[Size];
	 for(int i = 0;i<Size;i++) mpVisited[i]=false;
	 mMaxVisited = Size;
	 mStackPtr = -1;
	 mSize = 0;
  };

  ~STreeSearchStack(){delete mpStackData;delete mpVisited;};

  void visit(long Node){
  	if (Node<mMaxVisited){
  		mpVisited[Node] = true;
  	}else{
  		cout<<"Error. "<<Node<<" is out of bounds for visit.\n";
  	}
  }
  
  bool visited(long Node){return(mpVisited[Node]);};

  void clearVisited(){
	 for(int i = 0;i<mMaxVisited;i++){
		mpVisited[i]=false;
	 }
  }

  inline void push(long a){
	 mStackPtr++;
	 mpStackData[mStackPtr] = a;
	 mSize++;
  };

  inline long pop(){
	 mStackPtr--;
	 mSize--;
	 return(mpStackData[mStackPtr+1]);
  };

  inline int size(){return(mSize);};

};


//-------------------------------------------------------------------
// * SuffixTree
//------------------------------------------------------------------
// Class to represent a suffix tree for an arbitrary sequence.  Includes functionality
// for creating the suffix tree and searching the suffix tree.   Suffix trees can 
// be created in O(n) time.   Once constructed, searching a suffix tree takes time O(m), where 
// m is the size of the -query-.   Thus, suffix tree searching is very fast.  The 
// tradeoff is that suffix trees are fairly space inefficient.   If the queries are all 
// known to be of a fixed size, for example k-mers, then a hash will work better since
// it will have average case constant search time, with O(n) query time.  Moreover, it is 
// easier to optimize the space usage of a hash.  Suffix trees are superior if the 
// query size is not a constant.   Suffix trees are also the starting point for several 
// other interesting algorithms which, while not yet implemented here, form part of 
// the motivation for writing this class. 
//
class SuffixTree{
public:

  SuffixTree(){
	 Initialize();
  };
  
  ~SuffixTree(){
    Clear();
	 delete mpStack;
  }

  void Initialize(){
    mInvalidEdge.Invalidate();
    mNodeCount = 1;  // 1 because root node is defined
    mLeafNum = 0;
  }
	
  long mNodeCount;
  long mLeafNum;
  SuffixTreeNode *mNodes;
  STreeSearchStack *mpStack;

  void Clear(){
	 // Remove the nodes
    if (mNodes != NULL){
      delete [] mNodes;
    }
	 // Remove the source string.
	 if (mSource !=NULL){
		delete mSource;
	 }
  }

  // The input buffer and character count.  Please note that N
  // is the length of the input string -1, which means it
  // denotes the maximum index in the input buffer.
  const char *mSource;
  long mMaxSourceIndex;
	
  // KJD Need a version that counts Ns as match.
  void Create(const char *source); 
  long Find(const char *query,long length);
  long FindFirst(const char *query,long length,long &Location);
  long FindPrefix(const char *query,long length,long &Location);
  void FindAll(const char *query,long length,vector<long> &Locations);
  void PrintEdges();

  // protected:
	
  SuffixTreeEdge mInvalidEdge;  // Convenience object. 
  
  bool IsBase(char a);  
  inline bool EdgeValid(SuffixTreeEdge *pEdge);
  inline bool NodeIDValid(long NodeID);
  bool EdgeIDValid(int ID);
  inline void StackChildren(long NodeID,STreeSearchStack *pStack,SuffixTreeEdge *pTemp);
  void DFSLeafIdx(long NodeID,vector<long>&Indexes);
  inline SuffixTreeEdge * GetEdge(long nodeID,char edgeChar);
  inline SuffixTreeEdge * GetEdgeWithID(long nodeID,int edgeID);
  inline int CharToIndex(char);
  void CanonizeSuffix(SuffixTreeSuffix *pSuffix);
  void AddNewLeafEdge(long newFirstChar,long newLastChar,
		      long newStartNode,long);
  inline long SplitEdge(SuffixTreeEdge *pEdgeIn,SuffixTreeSuffix &s );
  inline bool MatchingEdgeExists(SuffixTreeSuffix &active,int LastCharIdx,SuffixTreeEdge* &pEdge);
  void AddPrefix(SuffixTreeSuffix &,long);
};


//-------------------------------------------------------------------
// * SuffixTreeSearchEngine
//------------------------------------------------------------------
// Wrapper for SuffixTree to adher to SequenceSeqrchEngine interface as well as handle 
// various practical details, such as splitting up large databases into several smaller 
// queries in order to stay within memory bounds. 
//
//class SuffixTreeSearchEngine: public SuffixTree, public SequenceSearchEngine{
class SuffixTreeSearchEngine: public SuffixTree{
public:

  string mSourceName;
  long   mSourceSize;

  SuffixTreeSearchEngine(){};
  //~SuffixTreeSearchEngine();

  // This causes a core dump on Linux, but no problem on solaris.... 
  //~SuffixTreeSearchEngine(){RemoveST();};  

  void CreateST(Sequence &Sequence);
  void Search(SequenceVector &Queries,SearchResultVec &Results);
  void RemoveST();

  void SearchAgainstDatabase(SequenceVector &Queries,string &DatabaseName,
			     SearchResultVec &Results);

};

//-------------------------------------------------------------------
// * SuffixTreeSearchEngineConcatenated
//------------------------------------------------------------------
// This idea only works if sources are the same size
//
//
class SuffixTreeSearchEngineConcatenated: public SuffixTreeSearchEngine{
public:

  Sequence mConcatenatedSource;
  vector<long> mFragmentBoundaries;  
  bool mbFilterSyntheticHits;

  long mSyntheticHitCount;
  long mSubjectsStartIdx;

  void ExtractSourceInfo(long cIdx,int &sIdx,long &sLoc);
  bool Synthetic(long loc,int qsize);
  void CreateST(SequenceVector &SeqVec);
  void CreateST(SequenceVector &SeqVec,int sstart,int send);
  void Search(SequenceVector &Queries, SearchResultVec &Results);
};

#endif










